#!/usr/bin/env python


def compute_Wt(num, model_adjustments, invar):
    """ Method of computing model error variance based on most recent
    N values.  If fewer than N values are available, then return a
    pre-specified constant (Galanis 2002) """
    import numpy as np

    Wt_init  = 0.07    # The default value
    num = num - 1
    length_used = 7

    if num < length_used:
        Wt = Wt_init
        return Wt
    else:
        index_start = num - length_used
        # Take the most recent model adjustments
        adjs = model_adjustments[index_start:num]
        if invar == 'precip':
            mean_std = np.std(adjs)
            Wt = float(np.random.normal(0.,mean_std,1))
        else:
            mean_var = np.var(adjs)
            Wt = mean_var / (length_used - 1)

        return Wt / 6.

def compute_Vt(num, model_adjustments, obs_adjustments, invar):
    """ Method for computing variance of obs equation
    based on the last N values (N days) of obs (Galanis 2002) """
    import numpy as np

    Vt_init = 0.4

    # How fast the filter changes depends on W/V (Homleid 1995)
    # Happy medium value of W/V = 0.06
    num = num - 1
    length_used = 7
    if num < length_used:
        Vt = Vt_init
        return Vt
    else:
        index_start = num - length_used
        # Grab the appropriate model adjustments
        adjs = obs_adjustments[index_start:num]
        if invar == 'precip':
            mean_std = np.std(adjs)
            Vt = float(np.random.normal(0.,mean_std,1))
        else:
            mean_var = np.var(adjs)
            Vt = mean_var / (length_used - 1)
        return Vt/6.




def run_filter(truth_dict, model_dict, maxlen, invar):
    """ Run a Kalman filter and return the time series of Kalman-filter
    corrected errors"""
    import numpy as np

    # Build the date lists
    model_dates = model_dict.keys()
    model_dates.sort()
    model_dates = model_dates[-maxlen:]

    truth_dates = truth_dict.keys()
    truth_dates.sort()
    truth_dates = truth_dates[truth_dates.index(model_dates[0]):]

    # Now build the plotting list for each variable
    model_vals = [model_dict[d][invar] for d in model_dates]
    truth_vals = [truth_dict[d][invar] for d in truth_dates]



    # Some initial values
    y_t = 0.
    P_t = 0.
    x_t = 0.
    initial_wt = 0

    model_adjustments = []
    obs_adjustments = []
    forecasts = [np.nan] # First day there is no forecast
    fdates = [model_dates[0]]
    P_list = [np.nan] # No covariance on first day either

    for n in range(len(truth_vals)):
        truth = truth_vals[n]
        model = model_vals[n]
        curdate = model_dates[n]

        # Make the current values be the previous values
        # This is our forward model--simple persistence
        # We'll then perturb these estimates in accordance with
        # an evolving model error
        y_prev = y_t   # The error observed
        x_prev = x_t   # The error estimate
        P_prev = P_t   # The covariance

        # Make a new estimate of the current error by taking the
        # prior_estimate (x_prev) and adding random error
        Wt = compute_Wt(n, model_adjustments, invar)
        x_est = x_prev + Wt

        # Now get the actual observed error at this time
        y_t = (model-truth)

        # Compute Vt (the observation error covariance)
        Vt = compute_Vt(n, model_adjustments, obs_adjustments, invar)

        # Estimate of Pt--the new covariance matrix.  Have to add in
        # random model error
        P_est = P_prev + Wt

        # Now compute the Kalman Gain.  Since we are letting Wt and Vt
        # evolve, this will be different each time
        Kt = P_est/(P_est + Vt)

        # Use the Kalman gain to improve the estimate of x_t based on
        # our observed error (y_t)
        x_t = x_est + Kt * (y_t - x_est)

        # For record keeping, keep a log of these
        model_adjustments.append(x_t)
        obs_adjustments.append(y_t - x_est)

        # Update the covariance with the Kalman gain
        P_t = (1-Kt)*P_est

        # Now make a corrected model forecast by adding our updated error (x_t) to the next
        # model forecast
        x_fore = model_vals[n+1] - x_t
        forecasts.append(x_fore)
        P_list.append(P_t)
        fdates.append(model_dates[n+1])

    return fdates, forecasts

        
